import yfinance as yf
import pandas as pd
from datetime import datetime
from typing import Union, Optional, Dict


class StockDataFetcher:
    """
    股票数据获取工具类（支持多时间粒度、多指标下载）

    功能特性：
    - 支持港股/美股/A股代码自动识别
    - 提供历史数据和实时数据获取
    - 自动处理时区转换
    - 内置数据缓存机制
    """

    def __init__(self, proxy: Optional[str] = None):
        """
        初始化下载器

        参数:
            proxy: 代理服务器地址（如"http://127.0.0.1:1080"）
        """
        if proxy:
            yf.set_proxy(proxy)

        # 市场代码映射表
        self.market_prefix = {
            'HK': '.HK',  # 港股
            'SS': '.SS',  # 沪股
            'SZ': '.SZ',  # 深股
            'US': ''  # 美股
        }

    def _format_ticker(self, ticker: str) -> str:
        """标准化股票代码格式"""
        if not any(ticker.endswith(suffix) for suffix in self.market_prefix.values()):
            if ticker[-2:].isdigit():  # 港股代码
                return ticker + self.market_prefix['HK']
            elif ticker.startswith('6'):
                return ticker + self.market_prefix['SS']
            elif ticker.startswith(('0', '3')):
                return ticker + self.market_prefix['SZ']
        return ticker

    def get_historical_data(
            self,
            ticker: str,
            start_date: Union[str, datetime],
            end_date: Union[str, datetime],
            interval: str = '1d',
            metrics: list = ['Open', 'High', 'Low', 'Close', 'Volume'],
            auto_adjust: bool = True
    ) -> pd.DataFrame:
        """
        获取历史行情数据

        参数:
            ticker: 股票代码（支持带市场后缀或自动识别）
            start_date: 开始日期（格式：YYYY-MM-DD）
            end_date: 结束日期
            interval: 数据间隔（1m/5m/15m/30m/60m/1d/1wk/1mo）
            metrics: 需要获取的指标列表
            auto_adjust: 是否自动调整价格（除权除息）

        返回:
            包含指定指标的DataFrame（索引为日期）

        示例:
            # >>> df = fetcher.get_historical_data('00700.HK', '2024-01-01', '2024-12-31')
            # >>> df = fetcher.get_historical_data('AAPL', '2024-01-01', interval='1h')
        """
        try:
            # formatted_ticker = self._format_ticker(ticker)
            data = yf.download(
                tickers=ticker,
                start=start_date,
                end=end_date,
                interval=interval,
                auto_adjust=auto_adjust,
                threads=True
            )

            # 筛选指定指标
            valid_metrics = [col for col in metrics if col in data.columns]
            return data[valid_metrics] if valid_metrics else data

        except Exception as e:
            raise ValueError(f"获取{ticker}历史数据失败: {str(e)}")

    def get_realtime_quote(
            self,
            ticker: str,
            fields: list = ['price', 'volume', 'marketCap']
    ) -> Dict:
        """
        获取实时行情快照

        参数:
            ticker: 股票代码
            fields: 需要获取的字段

        返回:
            包含实时数据的字典

        示例:
            # >>> quote = fetcher.get_realtime_quote('01810.HK')
        """
        try:
            stock = yf.Ticker(self._format_ticker(ticker))
            info = stock.fast_info

            return {
                field: getattr(info, field, None)
                for field in fields if hasattr(info, field)
            }

        except Exception as e:
            raise ValueError(f"获取{ticker}实时报价失败: {str(e)}")

    def get_dividends(self, ticker: str) -> pd.DataFrame:
        """获取分红派息历史"""
        try:
            stock = yf.Ticker(self._format_ticker(ticker))
            return stock.dividends.to_frame(name='Dividend')
        except Exception as e:
            raise ValueError(f"获取{ticker}分红数据失败: {str(e)}")


# 使用示例
if __name__ == "__main__":
    fetcher = StockDataFetcher()

    # 示例1：获取小米集团(01810.HK)2024年日线数据
    mi_history = fetcher.get_historical_data(
        ticker='002165',
        start_date='2025-01-20',
        end_date='2025-04-10',
        interval='1d'
    )

    mi_history.to_csv('./data/stock_data/hk01810_stock_data.csv')
    print(mi_history.head())

    # 示例2：获取苹果公司(AAPL)实时报价
    # aapl_quote = fetcher.get_realtime_quote('AAPL')
    # print(aapl_quote)